{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:127: RuntimeWarning: overflow encountered in add\n",
      "  a = a + a\n",
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:129: RuntimeWarning: invalid value encountered in subtract\n",
      "  temp1 = temp - a\n",
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:138: RuntimeWarning: invalid value encountered in subtract\n",
      "  itemp = int_conv(temp-a)\n",
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:162: RuntimeWarning: overflow encountered in add\n",
      "  a = a + a\n",
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:164: RuntimeWarning: invalid value encountered in subtract\n",
      "  temp1 = temp - a\n",
      "/home/leeyh/anaconda2/lib/python2.7/site-packages/numpy/core/machar.py:171: RuntimeWarning: invalid value encountered in subtract\n",
      "  if any(temp-a != zero):\n"
     ]
    }
   ],
   "source": [
    "import h5py\n",
    "import scipy.io as io\n",
    "import PIL.Image as Image\n",
    "import numpy as np\n",
    "import os\n",
    "import glob\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.ndimage.filters import gaussian_filter \n",
    "import scipy\n",
    "import json\n",
    "from matplotlib import cm as CM\n",
    "from image import *\n",
    "from model import CSRNet\n",
    "import torch\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#this is borrowed from https://github.com/davideverona/deep-crowd-counting_crowdnet\n",
    "def gaussian_filter_density(gt):\n",
    "    print gt.shape\n",
    "    density = np.zeros(gt.shape, dtype=np.float32)\n",
    "    gt_count = np.count_nonzero(gt)\n",
    "    if gt_count == 0:\n",
    "        return density\n",
    "\n",
    "    pts = np.array(zip(np.nonzero(gt)[1], np.nonzero(gt)[0]))\n",
    "    leafsize = 2048\n",
    "    # build kdtree\n",
    "    tree = scipy.spatial.KDTree(pts.copy(), leafsize=leafsize)\n",
    "    # query kdtree\n",
    "    distances, locations = tree.query(pts, k=4)\n",
    "\n",
    "    print 'generate density...'\n",
    "    for i, pt in enumerate(pts):\n",
    "        pt2d = np.zeros(gt.shape, dtype=np.float32)\n",
    "        pt2d[pt[1],pt[0]] = 1.\n",
    "        if gt_count > 1:\n",
    "            sigma = (distances[i][1]+distances[i][2]+distances[i][3])*0.1\n",
    "        else:\n",
    "            sigma = np.average(np.array(gt.shape))/2./2. #case: 1 point\n",
    "        density += scipy.ndimage.filters.gaussian_filter(pt2d, sigma, mode='constant')\n",
    "    print 'done.'\n",
    "    return density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#set the root to the Shanghai dataset you download\n",
    "root = '/home/leeyh/Downloads/Shanghai/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#now generate the ShanghaiA's ground truth\n",
    "part_A_train = os.path.join(root,'part_A_final/train_data','images')\n",
    "part_A_test = os.path.join(root,'part_A_final/test_data','images')\n",
    "part_B_train = os.path.join(root,'part_B_final/train_data','images')\n",
    "part_B_test = os.path.join(root,'part_B_final/test_data','images')\n",
    "path_sets = [part_A_train,part_A_test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img_paths = []\n",
    "for path in path_sets:\n",
    "    for img_path in glob.glob(os.path.join(path, '*.jpg')):\n",
    "        img_paths.append(img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for img_path in img_paths:\n",
    "    print img_path\n",
    "    mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))\n",
    "    img= plt.imread(img_path)\n",
    "    k = np.zeros((img.shape[0],img.shape[1]))\n",
    "    gt = mat[\"image_info\"][0,0][0,0][0]\n",
    "    for i in range(0,len(gt)):\n",
    "        if int(gt[i][1])<img.shape[0] and int(gt[i][0])<img.shape[1]:\n",
    "            k[int(gt[i][1]),int(gt[i][0])]=1\n",
    "    k = gaussian_filter_density(k)\n",
    "    with h5py.File(img_path.replace('.jpg','.h5').replace('images','ground_truth'), 'w') as hf:\n",
    "            hf['density'] = k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#now see a sample from ShanghaiA\n",
    "plt.imshow(Image.open(img_paths[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "gt_file = h5py.File(img_paths[0].replace('.jpg','.h5').replace('images','ground_truth'),'r')\n",
    "groundtruth = np.asarray(gt_file['density'])\n",
    "plt.imshow(groundtruth,cmap=CM.jet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.sum(groundtruth)# don't mind this slight variation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#now generate the ShanghaiB's ground truth\n",
    "path_sets = [part_B_train,part_B_test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img_paths = []\n",
    "for path in path_sets:\n",
    "    for img_path in glob.glob(os.path.join(path, '*.jpg')):\n",
    "        img_paths.append(img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for img_path in img_paths:\n",
    "    print img_path\n",
    "    mat = io.loadmat(img_path.replace('.jpg','.mat').replace('images','ground_truth').replace('IMG_','GT_IMG_'))\n",
    "    img= plt.imread(img_path)\n",
    "    k = np.zeros((img.shape[0],img.shape[1]))\n",
    "    gt = mat[\"image_info\"][0,0][0,0][0]\n",
    "    for i in range(0,len(gt)):\n",
    "        if int(gt[i][1])<img.shape[0] and int(gt[i][0])<img.shape[1]:\n",
    "            k[int(gt[i][1]),int(gt[i][0])]=1\n",
    "    k = gaussian_filter(k,15)\n",
    "    with h5py.File(img_path.replace('.jpg','.h5').replace('images','ground_truth'), 'w') as hf:\n",
    "            hf['density'] = k"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
